/*****************************************************************************************************
USAGE:
./opencv_annotation -images <folder location> -annotations <ouput file>

Created by: Puttemans Steven
*****************************************************************************************************/

#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>

#include <fstream>
#include <iostream>

using namespace std;
using namespace cv;

#include "libs/path/path_helper.h"

// default maximum size to show image
const int g_default_max_image_width = 640;
const int g_default_max_image_height = 640;

// Function prototypes
void on_mouse(int, int, int, int, void*);
string int2string(int);
void get_annotations(Mat, stringstream*);

// Public parameters
Mat image;
int roi_x0 = 0, roi_y0 = 0, roi_x1 = 0, roi_y1 = 0, num_of_rec = 0;
bool start_draw = false;

// Window name for visualisation purposes
const string window_name = "OpenCV Based Annotation Tool";

// FUNCTION : Mouse response for selecting objects in images
// If left button is clicked, start drawing a rectangle as long as mouse moves
// Stop drawing once a new left click is detected by the on_mouse function
void on_mouse(int event, int x, int y, int, void *)
{
	// Action when left button is clicked
	if (event == CV_EVENT_LBUTTONDOWN)
	{
		if (!start_draw)
		{
			roi_x0 = x;
			roi_y0 = y;
			start_draw = true;
		}
		else {
			roi_x1 = x;
			roi_y1 = y;
			start_draw = false;
		}
	}
	// Action when mouse is moving
	if ((event == CV_EVENT_MOUSEMOVE) && start_draw)
	{
		// Redraw bounding box for annotation
		Mat current_view;
		image.copyTo(current_view);
		rectangle(current_view, Point(roi_x0, roi_y0), Point(x, y), Scalar(0, 0, 255));
		imshow(window_name, current_view);
	}
}

// FUNCTION : snippet to convert an integer value to a string using a clean function
// instead of creating a stringstream each time inside the main code
string int2string(int num)
{
	stringstream temp_stream;
	temp_stream << num;
	return temp_stream.str();
}

// FUNCTION : given an image containing positive object instances, add all the object
// annotations to a known stringstream
void get_annotations(Mat input_image, stringstream* output_stream)
{
	// Make it possible to exit the annotation
	bool stop = false;

	// Reset the num_of_rec element at each iteration
	// Make sure the global image is set to the current image
	num_of_rec = 0;
	image = input_image;

	// Init window interface and couple mouse actions
	namedWindow(window_name, WINDOW_AUTOSIZE);
	setMouseCallback(window_name, on_mouse);

	imshow(window_name, image);
	stringstream temp_stream;
	int key_pressed = 0;

	do
	{
		// Keys for processing
		// You need to select one for confirming a selection and one to continue to the next image
		// Based on the universal ASCII code of the keystroke: http://www.asciitable.com/
		//      c = 99		    add rectangle to current image
		//	    n = 110		    save added rectangles and show next image
		//	    <ESC> = 27      exit program
		key_pressed = 0xFF & waitKey(0);
		switch (key_pressed)
		{
		case 27:
			destroyWindow(window_name);
			stop = true;
		case 99:
			// Add a rectangle to the list
			num_of_rec++;
			// Draw initiated from top left corner
			if (roi_x0 < roi_x1 && roi_y0<roi_y1)
			{
				temp_stream << " " << int2string(roi_x0) << " " << int2string(roi_y0) << " " << int2string(roi_x1 - roi_x0) << " " << int2string(roi_y1 - roi_y0);
			}
			// Draw initiated from bottom right corner
			if (roi_x0>roi_x1 && roi_y0 > roi_y1)
			{
				temp_stream << " " << int2string(roi_x1) << " " << int2string(roi_y1) << " " << int2string(roi_x0 - roi_x1) << " " << int2string(roi_y0 - roi_y1);
			}
			// Draw initiated from top right corner
			if (roi_x0 > roi_x1 && roi_y0 < roi_y1)
			{
				temp_stream << " " << int2string(roi_x1) << " " << int2string(roi_y0) << " " << int2string(roi_x0 - roi_x1) << " " << int2string(roi_y1 - roi_y0);
			}
			// Draw initiated from bottom left corner
			if (roi_x0<roi_x1 && roi_y0>roi_y1)
			{
				temp_stream << " " << int2string(roi_x0) << " " << int2string(roi_y1) << " " << int2string(roi_x1 - roi_x0) << " " << int2string(roi_y0 - roi_y1);
			}
			temp_stream << endl;

			rectangle(input_image, Point(roi_x0, roi_y0), Point(roi_x1, roi_y1), Scalar(0, 255, 0), 1);

			break;
		}

		// Check if escape has been pressed
		if (stop)
		{
			break;
		}
	}
	// Continue as long as the next image key has not been pressed
	while (key_pressed != 110);

	// If there are annotations AND the next image key is pressed
	// Write the image annotations to the file
	if (num_of_rec > 0 && key_pressed == 110)
	{
		*output_stream << " " << num_of_rec << endl << temp_stream.str() << endl;
	}

	// Close down the window
	destroyWindow(window_name);
}

void load_annotations(const string& annotation_filename, vector<Rect>& annotations)
{
	ifstream infile(annotation_filename);
	if (infile.is_open())
	{
		int n;
		infile >> n;
		annotations.resize(n);
		for (size_t i = 0; i < n; i++)
		{
			infile >> annotations[i].x >> annotations[i].y >> annotations[i].width >> annotations[i].height;
		}
		infile.close();
	}
}

void save_annotations(const string& annotation_filename, const vector<Rect>& annotations)
{
	ofstream outfile(annotation_filename);
	if (outfile.is_open())
	{
		outfile << annotations.size() << endl;
		for (size_t i = 0; i < annotations.size(); i++)
		{
			outfile << annotations[i].x << " " << annotations[i].y << " " << annotations[i].width << " " << annotations[i].height << endl;
		}
		outfile.close();
	}
}

int main(int argc, const char** argv)
{
	// Read in the input arguments
	string image_folder;
	string annotation_folder;
	int flag = 0;
	for (int i = 1; i < argc; ++i)
	{
		if (!strcmp(argv[i], "-images"))
		{
			image_folder = argv[++i];
			flag++;
		}
		else if (!strcmp(argv[i], "-annotations"))
		{
			annotation_folder = argv[++i];
			flag++;
		}
	}

	if (flag != 2)
	{
		cerr << "usage:" << argv[0] << " -images <images path> -annotations <annotations path>" << endl;
		return -1;
	}

	// Return the image filenames inside the image folder
	vector<String> filenames;
	JDir dir;
	dir.clear_suffix().add_suffix(".jpg").add_suffix(".png").add_suffix(".bmp")(image_folder, filenames);
	

	double scale = 1.0;
	bool isScale = false;
	// Loop through each image stored in the images folder
	// Create and temporarily store the annotations
	// At the end write everything to the annotations file
	for (size_t i = 0; i < filenames.size(); i++){
		string filename = filenames[i];
		// Read in an image
		Mat current_image = imread(image_folder + "/" + filename);

		if (current_image.rows > g_default_max_image_height)
		{
			isScale = true;
			scale = 1.0 * g_default_max_image_height / current_image.rows;
			int n_width = int(current_image.cols * scale + 0.5);
			int n_height = int(current_image.rows * scale + 0.5);
			Mat tmp_image;
			resize(current_image, tmp_image, Size(n_width, n_height));
			current_image = tmp_image;
		}

		string annotaions_filename = filename.substr(0, filename.find_last_of(".")) + ".objs";
		//cout << filename << " " << annotaions_filename << endl;
		vector<Rect> annotations;
		load_annotations(annotation_folder + "/" + annotaions_filename, annotations);
		cout << "loaded " << annotations.size() << " rectangles." << endl;

		for (size_t j = 0; j < annotations.size(); j++)
		{
			Rect tmp_rect = annotations[j];
			if (isScale)
			{
				tmp_rect.x *= scale;
				tmp_rect.y *= scale;
				tmp_rect.width *= scale;
				tmp_rect.height *= scale;
			}
			rectangle(current_image, tmp_rect, Scalar(0, 0, 255), 2);
		}
		// Perform annotations & generate corresponding output
		stringstream output_stream;
		get_annotations(current_image, &output_stream);

		// Store the annotations, write to the output file
		if (output_stream.str() != ""){
			int n, x, y, w, h;
			output_stream >> n;
			for (size_t j = 0; j < n; j++)
			{
				output_stream >> x >> y >> w >> h;
				if (isScale)
				{
					annotations.push_back(Rect(x / scale, y / scale, w / scale, h / scale));
				}
				else
				{
					annotations.push_back(Rect(x, y, w, h));
				}
				
			}
		}

		save_annotations(annotation_folder + "/" + annotaions_filename, annotations);
		isScale = false;
	}

	return 0;
}
